import tensorflow.contrib.keras as kr
import tensorflow as tf
import pickle
import numpy as np
import time

test_content_list = ["反腐败斗争取得压倒性胜利”。这是2018年12月13日召开的中央政治局会议对我国反腐败斗争形势的最新重大判断。2017年10月，党的十九大对反腐败斗争形势作出的判断是——“反腐败斗争压倒性态势已经形成并巩固发展”。从形成“压倒性态势”到取得“压倒性胜利”，标志着我国反腐败斗争成果正从量的积累迈向质的转变。"]
sequence_length = 600
#3-加载词汇表
with open('vocabulary_list.pickle', 'rb') as file:
    vocabulary_list = pickle.load(file)
    word2id_dict = dict([(b, a) for a, b in enumerate(vocabulary_list)])
    content2idList = lambda content : [word2id_dict[word] for word in content if word in word2id_dict]
    test_idlist_list = [content2idList(content) for content in test_content_list]
    test_X = kr.preprocessing.sequence.pad_sequences(test_idlist_list, sequence_length)
    #print("test_X:", test_X)

with open('label_list.pickle', 'rb') as file:
    label_list = pickle.load(file)
    num_classes = np.unique(label_list)
    print("num_classes:", num_classes)

t1 = time.time()
with tf.Session() as session:
    saver = tf.train.import_meta_graph("save/my-model-10000.meta")
    saver.restore(session, tf.train.latest_checkpoint("save/"))
    predict_Y = tf.get_collection('pred_network')[0]
    graph = tf.get_default_graph()
    X_holder = graph.get_operation_by_name('X_holder').outputs[0]
    keep_prob = graph.get_operation_by_name('keep_prob').outputs[0]
    predict_value = session.run(predict_Y, feed_dict={X_holder: test_X, keep_prob: 1.0})
    Y = np.array(predict_value[0])
    print("这是", num_classes[np.argmax(Y)], "新闻")
t2 = time.time()
print("time:",t2-t1)